"""
数据基类， 为对接choices量化数据接口
"""
import datetime
import os
from typing import Iterable

import numpy as np
import pandas as pd
from EmQuantAPI import *

from dytz.bar import Bar

__all__ = ['ChoicesData', 'TimeSeries']


class ChoicesData:
    def __init__(self, code: str, startdate: str, enddate: str,
                 ind: str = "OPEN,CLOSE,HIGH,LOW,VOLUME,AVGPRICE", data_org: str = 'locality'):
        """
        构造器
        :param code: 回测标的
        :param startdate: 开始时间
        :param enddate: 结束时间
        :param ind: 数据字段
        :param data_org: 数据来源， choices: choices获取， locality: 本地，需要事先下载
        """
        if data_org == 'choices':  # 从choices获取数据

            if os.path.exists(f'date_{startdate}_{enddate}.pkl'):
                self.date = pd.read_pickle(f'date_{startdate}_{enddate}.pkl')
            else:
                self.loginResult = c.start("ForceLogin=1")
                self.date = pd.Series(pd.to_datetime(c.tradedates(startdate, enddate).Data))
                self.date.to_pickle(f'date_{startdate}_{enddate}.pkl')

            self.T = len(self.date)
            self.t = 0
            if os.path.exists(f'{code.replace(".SH", "")}_day.pkl'):
                data = pd.read_pickle(f'{code.replace(".SH", "")}_day.pkl')
                data.index = self.date
                data.loc[:, 'TICKER'] = code
                if 'AVGPRICE' not in data.columns:
                    data.loc[:, 'AVGPRICE'] = np.nan
                self.data = data
            else:
                data = c.csd(code, ind, startdate,
                             enddate, "period=1,adjustflag=3,curtype=2,order=1,market=CNSESH").Data
                ind_ = ind.split(',')
                self.data = pd.DataFrame()
                code = code.split(',')
                for co in code:
                    da = pd.DataFrame(data[co], index=ind_, columns=self.date).T
                    da.loc[:, 'TICKER'] = co
                    self.data = pd.concat([self.data, da], axis=0)
        elif data_org == 'locality':
            #  本地数据
            data = pd.read_excel(f'K线导出_{code.replace(".SH", "")}_日线数据.xls')
            data.columns = ['code', 'name', 'date', 'OPEN', 'HIGH', 'LOW', 'CLOSE', 'a', 'b', 'VOLUME', 'm']
            data.loc[:, 'AVGPRICE'] = np.nan
            data = data.loc[:, ['date'] + ind.split(',')]
            data.loc[:, 'date'] = pd.to_datetime(data.loc[:, 'date'])
            data = data[(data.date >= datetime.datetime.strptime(startdate, '%Y%m%d')) &
                        (data.date <= datetime.datetime.strptime(enddate, '%Y%m%d'))]
            self.date = data.loc[:, 'date']
            self.data = data.set_index('date')
            self.data.loc[:, 'TICKER'] = code
            self.t = 0
            self.T = self.data.shape[0]
        else:
            raise "未指定数据获取类型"

    def __iter__(self):
        return self

    def __next__(self):
        if self.t < self.T:
            bar = Bar(self.data.iloc[self.t])
            self.t += 1
            return bar
        else:
            raise StopIteration


class TimeSeries:
    """时间序列基类"""

    def __init__(self, n: int, dtype: type = float, default: int or float or Iterable = 0):
        """
        构造器
        :param n: 序列长度
        :param dtype: 序列类型
        :param default: 默认值，iter必须和n相等
        """
        if isinstance(default, Iterable):
            self.timearray = np.asarray(default)
        else:
            self.timearray = np.asarray([dtype(default)] * int(n))

    def set(self, x: float or int) -> None:
        """
        将序列的最后一个值改编为x
        :param x: 需要改变的值
        :return: None
        """
        self.timearray[-1] = x

    def step(self) -> None:
        """
        将序列迁移一位
        :return: None
        """
        self.timearray[0: -1] = self.timearray[1:]
        self.set(self.timearray[-1])

    def __repr__(self):
        return repr(self.timearray)

    def __getitem__(self, item):
        if isinstance(item, slice):
            return self.timearray[::-1][item][::-1]
        else:
            return self.timearray[::-1][item]

    def max(self):
        return np.max(self.timearray)

    def min(self):
        return np.min(self.timearray)

    def __eq__(self, other: str or int or np.array) -> np.array:
        """
        ==
        :param other: 可以为一个数，或者一个序列
        :return: 返回bool值的序列
        """
        return self.timearray == other

    def __lt__(self, other: str or int or np.array) -> np.array:
        """
        <
        :param other: 可以为一个数，或者一个序列
        :return: 返回bool值的序列
        """
        return self.timearray < other

    def __gt__(self, other: str or int or np.array) -> np.array:
        """
        >
        :param other: 可以为一个数，或者一个序列
        :return: 返回bool值的序列
        """
        return self.timearray > other

    def __le__(self, other: str or int or np.array) -> np.array:
        """
        <=
        :param other: 可以为一个数，或者一个序列
        :return: 返回bool值的序列
        """
        return self.timearray <= other

    def __ge__(self, other: str or int or np.array) -> np.array:
        """
        >=
        :param other: 可以为一个数，或者一个序列
        :return: 返回bool值的序列
        """
        return self.timearray >= other
